import warnings
warnings.filterwarnings("ignore")

from pathlib import Path
import sys
from urllib.error import HTTPError
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.pipeline import Pipeline
from sklearn.svm import LinearSVC
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt


def main():
    # 1) 加载数据（可选去噪）
    remove_parts = ("headers", "footers", "quotes")
    # 将数据缓存到项目根的 data_cache 目录，便于手动放置压缩包离线使用
    project_root = Path(__file__).resolve().parent.parent
    data_home = project_root / "data_cache"
    data_home.mkdir(parents=True, exist_ok=True)

    print(f"开始加载 20 Newsgroups 数据集...\n数据缓存目录: {data_home}")
    try:
        train = fetch_20newsgroups(subset="train", remove=remove_parts, data_home=str(data_home))
        test = fetch_20newsgroups(subset="test", remove=remove_parts, data_home=str(data_home))
    except HTTPError as e:
        # 403/网络受限时给出离线放置指引
        print("\n下载数据失败 (HTTPError)。请手动下载数据压缩包后重试：")
        sys.exit(1)

    X_train, y_train = train.data, train.target
    X_test, y_test = test.data, test.target
    target_names = train.target_names

    # 2) 建立 Pipeline（以 LinearSVC 为例）
    pipeline = Pipeline([
        ("tfidf", TfidfVectorizer(stop_words="english")),
        ("clf", LinearSVC())
    ])

    # 3) 快速版：直接训练，不做网格搜索
    print("开始训练快速模型（不进行网格搜索）...")
    pipeline.set_params(
        tfidf__max_df=0.75,
        tfidf__min_df=2,
        tfidf__ngram_range=(1, 2),
        clf__C=1.0
    )
    pipeline.fit(X_train, y_train)

    # 4) 测试集评估
    y_pred = pipeline.predict(X_test)
    print("测试集准确率:", accuracy_score(y_test, y_pred))
    print("分类报告:\n")
    print(classification_report(y_test, y_pred, target_names=target_names))

    # 5) 混淆矩阵可视化
    cm = confusion_matrix(y_test, y_pred)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, cmap="Blues", square=True, cbar=True)
    plt.title("混淆矩阵（LinearSVC）")
    plt.xlabel("预测类别")
    plt.ylabel("真实类别")
    plt.tight_layout()
    plt.show()

    # 6) 朴素贝叶斯对比
    print("\n开始训练朴素贝叶斯基线模型...")
    nb_pipeline = Pipeline([
        ("tfidf", TfidfVectorizer(stop_words="english")),
        ("clf", MultinomialNB())
    ])
    nb_pipeline.fit(X_train, y_train)
    nb_pred = nb_pipeline.predict(X_test)
    print("朴素贝叶斯 测试集准确率:", accuracy_score(y_test, nb_pred))


if __name__ == "__main__":
    main()


